#include "depth_to_space_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
namespace optiling {
    constexpr uint32_t ALIGN_NUM = 32;    
    constexpr uint32_t CALC_NUM = 256;  
    constexpr uint32_t BUFFER_NUM = 1;    
    uint32_t RoundUp(uint32_t a, uint32_t b)
    { 
        if (b == 0) {
            return 0;
        }
        return ((a + b - 1) / b) * b;
    }    
    uint32_t RoundDown(uint32_t a, uint32_t b)
    { 
        if (b == 0) {
            return 0;
        }        
        return (a / b) * b;
    }     
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    DepthToSpaceTilingData tiling;
    auto ascendcPlatform = platform_ascendc:: PlatformAscendC(context->GetPlatformInfo());
    auto socVersion = ascendcPlatform.GetSocVersion();
    if (!((socVersion == platform_ascendc::SocVersion::ASCEND910B))) {
        return ge::GRAPH_FAILED;
    }
    uint64_t ub_size;
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
    auto aivNum = ascendcPlatform.GetCoreNumAiv();
    const int64_t* pblock_size = context->GetAttrs()->GetInt(0);
    const char* pmode = context->GetAttrs()->GetStr(1); 
    const char* pdata_format = context->GetAttrs()->GetStr(2);
    const gert::StorageShape* x1_shape = context->GetInputShape(0);
    const gert::Shape shape = x1_shape->GetStorageShape();
    uint32_t totalLength = x1_shape->GetOriginShape().GetShapeSize();
    int64_t N, C, H, W, B;
    B = (* pblock_size);
    uint32_t coreNum;
    uint32_t WAlign;    
    uint32_t baseW;
    if((strcmp(pdata_format, "NCHW") == 0)){
        N = shape.GetDim(0);
        C = shape.GetDim(1);
        H = shape.GetDim(2);
        W = shape.GetDim(3);
    }
    else if((strcmp(pdata_format, "NHWC") == 0))
    {
        N = shape.GetDim(0);
        H = shape.GetDim(1);
        W = shape.GetDim(2);
        C = shape.GetDim(3);
    }
    else{
        N = shape.GetDim(0);
        H = shape.GetDim(1);
        W = shape.GetDim(2);
        C = shape.GetDim(3);
    }
    const gert::Tensor * x = context->GetInputTensor(0);
    uint32_t dTypeSize;
    auto x_dtype = x->GetDataType();  
    if((x_dtype == ge::DT_FLOAT16))
    {
        dTypeSize = 2; 
    }
    else if((x_dtype == ge::DT_FLOAT))
    {
        dTypeSize = 4; 
    } 
    else if((x_dtype == ge::DT_INT32))
    {
        dTypeSize = 4; 
    } 
    else if((x_dtype == ge::DT_INT8))
    {
        dTypeSize = 1; 
    }   
    else
    {
    }  
    uint32_t maxTileLength = (ub_size / BUFFER_NUM / ALIGN_NUM) * (ALIGN_NUM / dTypeSize);
    uint32_t elemPerBlock = ALIGN_NUM / dTypeSize;
    if((strcmp(pmode, "DCR") == 0) && (strcmp(pdata_format, "NHWC") == 0)) {
        if(W > aivNum) {
            coreNum = aivNum;  
        }
        else {
            coreNum = W;
        }
        WAlign = RoundUp(W, coreNum);    
        baseW = WAlign / coreNum;
        if(((C/B) % (elemPerBlock) == 0))
        { 
            context->SetTilingKey(1);
            if(baseW > (maxTileLength / C))
            {
                baseW = (maxTileLength / C);
            }
        }
        else
        {
            context->SetTilingKey(5);
            uint32_t C_B = C / B;
            uint32_t C_BAlign = RoundUp(C_B, elemPerBlock);
            if(baseW > (maxTileLength / (C_BAlign * B)))
            {
                baseW = (maxTileLength / (C_BAlign * B));
            }             
        }
        context->SetBlockDim(coreNum);  
        tiling.set_blockLength(totalLength);
        tiling.set_N(N);
        tiling.set_H(H);
        tiling.set_W(W);
        tiling.set_C(C);
        tiling.set_B(B);
        tiling.set_baseW(baseW);
    }
    if((strcmp(pmode, "DCR") == 0) && (strcmp(pdata_format, "NCHW") == 0)) { 
        if((W * RoundUp(B, elemPerBlock) + RoundUp((B * W), elemPerBlock)) > maxTileLength)
        {
        }
        else
        {
            context->SetTilingKey(2);
            uint32_t CBH = (C / B) * H;
            if(CBH > aivNum) {
                coreNum = aivNum;  
            }
            else {
                coreNum = CBH;
            }
            context->SetBlockDim(coreNum);  
            tiling.set_blockLength(totalLength);
            tiling.set_N(N);
            tiling.set_H(H);
            tiling.set_W(W);
            tiling.set_C(C);
            tiling.set_B(B);
        }      
    }
    if((strcmp(pmode, "CRD") == 0) && (strcmp(pdata_format, "NHWC") == 0)) {          
        if((2 * RoundUp((W * C), elemPerBlock)) > maxTileLength)
        {
        }   
        else
        {
            context->SetTilingKey(3);
            if(H > aivNum) {
                coreNum = aivNum;  
            }
            else {
                coreNum = H;
            }
            context->SetBlockDim(coreNum);  
            tiling.set_blockLength(totalLength);
            tiling.set_N(N);
            tiling.set_H(H);
            tiling.set_W(W);
            tiling.set_C(C);
            tiling.set_B(B);
        }  
    }
    if((strcmp(pmode, "CRD") == 0) && (strcmp(pdata_format, "NCHW") == 0)) {     
        if((W * RoundUp(B, elemPerBlock) + RoundUp((B * W), elemPerBlock)) > maxTileLength)
        {
        }
        else
        {             
            context->SetTilingKey(4);
            uint32_t CBH = (C / B) * H;
            if(CBH > aivNum) {
                coreNum = aivNum;  
            }
            else {
                coreNum = CBH;
            }
            context->SetBlockDim(coreNum);  
            tiling.set_blockLength(totalLength);
            tiling.set_N(N);
            tiling.set_H(H);
            tiling.set_W(W);
            tiling.set_C(C);
            tiling.set_B(B);
        }
    }
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    return ge::GRAPH_SUCCESS;
}
}
namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}
namespace ops {
class DepthToSpace : public OpDef {
public:
    explicit DepthToSpace(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("block_size").Int();
        this->Attr("mode").AttrType(OPTIONAL).String("None");
        this->Attr("data_format").AttrType(OPTIONAL).String("None");
        this->SetInferShape(ge::InferShape);
        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");
    }
};
OP_ADD(DepthToSpace);
}
